from functools import partial

import numpy

import sko
from scene.invest_opt.model.target import target_all


class INVEST_OPT:
    def __init__(self, args, opening_price):
        self.args = args
        self.opening_price = opening_price

        if self.args.method not in [0, 1, 2]:
            assert False, 'not vaild'

        # 设置随机数种子保证每次运行结果相同
        if args.seed is not None:
            numpy.random.seed(args.seed)

    def predict(self):

        costfun = partial(target_all, mat=self.opening_price, args=self.args)
        sko.tools.set_run_mode(costfun, 'vectorization')
        x0 = numpy.abs(numpy.random.rand(self.args.input_size + 1))

        if self.args.method == 0:
            opt = sko.SA.SA(
                func=costfun,
                x0=x0,
                L=self.args.L,
                T_min=self.args.Te,
                T_max=self.args.T0,
                max_stay_counter=self.args.epochs,
                lb=0,
                ub=1
            )
            self.ans, self.rec = opt.run()
        elif self.args.method == 1:
            opt = sko.PSO.PSO(
                func=costfun,
                n_dim=self.args.input_size + 1,
                pop=self.args.batch,
                max_iter=self.args.epochs,
                c1=self.args.c1,
                c2=self.args.c2,
                max_stay_iter=self.args.max_stay_step,
                lb=0,
                ub=1
            )
            self.ans, self.rec = opt.run()
        elif self.args.method == 2:

            opt = sko.DE.DE(
                func=costfun,
                n_dim=self.args.input_size + 1,
                size_pop=self.args.batch,
                max_iter=self.args.epochs,
                prob_mut=self.args.mut_std,
                F=self.args.mut_cof,
                max_stay_iter=self.args.max_stay_step,
                lb=0,
                ub=1
            )
            self.ans, self.rec = opt.run()
        else:
            assert False, 'not vaild'

    # 输出方案及其指标
    def getRec(self):
        return self.ans, self.rec
